package io.xccit.deepseek.controller;

import io.xccit.deepseek.enums.MessageType;
import io.xccit.deepseek.service.OllamaService;
import io.xccit.deepseek.utils.SseServer;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

/**
 * @author CH_ywx
 * @version 1.0
 * @create 2025/4/7
 * <p>SSE服务控制器</p>
 */
@RestController
@RequestMapping("/sse")
public class SseController {

    private final OllamaService ollamaService;

    @Autowired
    public SseController(OllamaService ollamaService) {
        this.ollamaService = ollamaService;
    }


    /**
     * 建立连接
     * @param userId
     * @return
     */
    @GetMapping(path="/connect", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter sseConnection(@RequestParam("userId") String userId) {
        return SseServer.sseConnection(userId);
    }

    /**
     * 发送消息
     * @param userId
     * @param message
     * @return
     */
    @GetMapping(path="/sendMessage",produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public String sendMessage(@RequestParam("userId") String userId,
                              @RequestParam("message") String message) {
        Flux<ChatResponse> chatResponseFlux = ollamaService.chatStream(message);
        chatResponseFlux.subscribe(chatResponse -> {
            String content = chatResponse.getResult().getOutput().getContent();
            SseServer.sendMessage(userId, content, MessageType.MESSAGE);
        });
        return "SUCCESS";
    }

    /**
     * 广播消息
     * @param message
     * @return
     */
    @GetMapping(path="/broadCast", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public String broadCast(@RequestParam("message") String message) {
        Flux<ChatResponse> chatResponseFlux = ollamaService.chatStream(message);
        chatResponseFlux.subscribe(chatResponse -> {
            String content = chatResponse.getResult().getOutput().getContent();
            SseServer.broadCast(content);
        });
        return "SUCCESS";
    }


    /**
     * 获取当前连接数
     * @return
     */
    @GetMapping("/count")
    public int getConnectedCount() {
        return SseServer.getCount();
    }

    /**
     * 断开连接
     * @param userId
     */
    @GetMapping("/stop")
    public String stopConnected(@RequestParam("userId") String userId) {
        SseServer.stopConnected(userId);
        return "SUCCESS";
    }
}
